from tiingo import TiingoClient
import pandas as pd
import numpy as np
from sklearn.preprocessing import RobustScaler, PowerTransformer, StandardScaler
from sklearn.model_selection import train_test_split
import os


class DataPipeline:
    def __init__(self, api_key=None):
        if not api_key:
            raise ValueError("未提供 Tiingo API 密钥。请通过参数提供")
        # 修改 TiingoClient 初始化方式
        self.client = TiingoClient({'api_key': api_key})

    def _fetch_raw_data(self):
        """获取原始市场数据"""
        return self.client.get_ticker_price(
            '600887',
            startDate=pd.Timestamp.now() - pd.DateOffset(years=3),
            endDate=pd.Timestamp.now(),
            frequency='daily',
            fmt='json'
        )

    def _process_raw_data(self, raw_data):
        df = pd.DataFrame(raw_data)[['date', 'open', 'high', 'low', 'close', 'volume']]
        df.columns = ['Date', 'Open', 'High', 'Low', 'Close', 'Volume']
        df['Date'] = pd.to_datetime(df['Date'])
        df.set_index('Date', inplace=True)
        # 异常值处理
        df = self._handle_outliers(df)
        # 特征工程
        df = self._add_features(df)
        return df.dropna()

    def _handle_outliers(self, df):
        # 价格特征
        price_cols = ['Open', 'High', 'Low', 'Close']
        for col in price_cols:
            q_low = df[col].quantile(0.05)
            q_high = df[col].quantile(0.95)
            df[col] = np.where(df[col] < q_low, q_low,
                               np.where(df[col] > q_high, q_high, df[col]))
        # 交易量特征（IQR方法）
        vol_q1 = df['Volume'].quantile(0.25)
        vol_q3 = df['Volume'].quantile(0.75)
        iqr = vol_q3 - vol_q1
        df['Volume'] = df['Volume'].clip(upper=vol_q3 + 3 * iqr)
        return df

    def _compute_macd(self, series, slow=26, fast=12, signal=9):
        ema_fast = series.ewm(span=fast).mean()
        ema_slow = series.ewm(span=slow).mean()
        macd = ema_fast - ema_slow
        signal_line = macd.ewm(span=signal).mean()
        return pd.DataFrame({'MACD': macd, 'Signal': signal_line}, index=series.index)

    def _compute_rsi(self, series, period=14):
        delta = series.diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(period).mean()
        avg_loss = loss.rolling(period).mean()
        rs = avg_gain / avg_loss
        return 100 - (100 / (1 + rs))

    def _compute_atr(self, df, period=14):
        high_low = df['High'] - df['Low']
        high_close = np.abs(df['High'] - df['Close'].shift())
        low_close = np.abs(df['Low'] - df['Close'].shift())
        tr = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
        return tr.rolling(period).mean()

    def _compute_obv(self, df):
        """能量潮指标"""
        obv = (np.sign(df['Close'].diff()) * df['Volume']).fillna(0).cumsum()
        return obv

    def _add_features(self, df):
        df = df.copy()

        # 基础特征
        df['Returns'] = df['Close'].pct_change()
        df['Volatility'] = df['Returns'].rolling(5).std()

        # 技术指标（确保索引对齐）
        df['RSI'] = self._compute_rsi(df['Close'])
        macd_df = self._compute_macd(df['Close'])
        df = pd.concat([df, macd_df], axis=1)  # 横向合并MACD指标

        df['ATR'] = self._compute_atr(df)
        df['OBV'] = self._compute_obv(df)

        # 滞后收益率特征
        for lag in [1, 3, 5]:
            df[f'Return_{lag}d'] = df['Returns'].shift(lag)

        return df.dropna()

    def get_processed_data(self):
        raw_data = self._fetch_raw_data()
        return self._process_raw_data(raw_data)


class ScalerPipeline:
    def __init__(self, train_data):
        self.scalers = {}
        self._fit_scalers(train_data)

    def _fit_scalers(self, data):
        """Fit scalers for different feature groups"""
        price_cols = ['Open', 'High', 'Low', 'Close']
        volume_cols = ['Volume']

        # Price scaler
        price_scaler = RobustScaler()
        price_scaler.fit(data[price_cols])
        self.scalers['price'] = price_scaler

        # Volume scaler
        volume_scaler = StandardScaler()
        volume_scaler.fit(data[volume_cols])
        self.scalers['volume'] = volume_scaler

    def transform(self, data):
        """Transform data using fitted scalers"""
        price_cols = ['Open', 'High', 'Low', 'Close']
        volume_cols = ['Volume']

        scaled_data = data.copy()
        scaled_data[price_cols] = self.scalers['price'].transform(data[price_cols])
        scaled_data[volume_cols] = self.scalers['volume'].transform(data[volume_cols])
        return scaled_data
